Conversation
Add `ns nemo_gym_rollouts` CLI command for orchestrating rollout collection with NeMo Gym, including self-hosted or pre-hosted vLLM servers, optional sandbox, and parallel seed-based scaling. New script classes: - NemoGymRolloutsScript: orchestrates ng_run + ng_collect_rollouts - MultiVLLMServerScript: data-parallel vLLM with multiple replicas per node - GymClientScript: multi-node vLLM server discovery via SLURM env vars Infrastructure enhancements: - Command: mounts, environment, workdir, avoid/force nemo_run_code controls - Per-script num_tasks_override for mixed task configurations - Allocation ordering fix for multi-component overlap jobs - NEMO_SKILLS_FORCE_PATTERN_PACKAGER env var for packaging control - --kill-on-bad-exit=1 srun flag for early failure detection Signed-off-by: George Armstrong <georgea@nvidia.com>
gwarmstrong
force-pushed
the
georgea/nemo-gym-rollouts-pr
branch
2 times, most recently
from
d61b5f4 to
d86e6c9
Compare
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| VLLM_SERVER_URL="{vllm_server_url}" | ||
| if [ -n "$VLLM_SERVER_URL" ]; then | ||
| echo "=== Waiting for vLLM server at $VLLM_SERVER_URL ===" | ||
| while [ $(curl -s -o /dev/null -w "%{{http_code}}" "$VLLM_SERVER_URL/models" 2>/dev/null) -ne 200 ]; do |
There was a problem hiding this comment.
The curl command returns the HTTP status code directly in the condition, but this can fail when curl itself fails (network issues, DNS failures, etc). The command substitution $(curl ...) will return an empty string on curl failure, which bash compares numerically to 200, causing an error.
| while [ $(curl -s -o /dev/null -w "%{{http_code}}" "$VLLM_SERVER_URL/models" 2>/dev/null) -ne 200 ]; do | |
| while ! curl -sf "$VLLM_SERVER_URL/models" >/dev/null 2>&1; do |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a Nemo Gym rollouts CLI and new script classes under a scripts subpackage; expands Command with mounts/environment/workdir and PYTHONPATH controls; introduces deterministic packaging flags, fail-fast Slurm srun, tests, docs, and removes the old monolithic scripts file. Changes
Sequence Diagram(s)sequenceDiagram
participant CLI as "User CLI\n(nemo_gym_rollouts)"
participant Orch as "Pipeline Orchestrator\n(Command / Pipeline)"
participant Server as "vLLM Server\n(ServerScript)"
participant Sandbox as "Sandbox\n(SandboxScript)"
participant Client as "Gym Client\n(NemoGymRolloutsScript)"
participant Storage as "Storage\n(output_dir / logs)"
CLI->>Orch: invoke nemo_gym_rollouts(args)
Orch->>Server: create/configure server jobs (ports, GPUs, nodes)
Orch->>Sandbox: optionally create sandbox per job
Orch->>Client: create client jobs with server addresses & env
Server-->>Client: provide server URL / health status
Client->>Server: run rollouts against vLLM
Client->>Storage: write collected rollouts and logs
Orch->>Storage: aggregate statuses and finalize
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Fix all issues with AI agents
In `@nemo_skills/pipeline/nemo_gym_rollouts.py`:
- Line 238: The code uses cluster_config["containers"].get(server_type_str,
server_type_str) which can silently accept typos as literal image names; change
this to validate the server_type_str against cluster_config["containers"] (or
use direct access) so missing keys fail fast: check if server_type_str is a key
in cluster_config["containers"] and if so set server_container =
cluster_config["containers"][server_type_str], otherwise raise a clear error (or
fallback to a validated explicit list) so functions referencing server_container
(and the server_type_str lookup) don’t silently accept invalid names.
In `@nemo_skills/pipeline/utils/scripts.py`:
- Around line 930-936: The error block that runs when BEST_COUNT is less than
NUM_NODES prints several echo lines with inconsistent indentation; normalize the
leading spaces so all echo lines align with the surrounding echo and exit
statements. Locate the conditional using BEST_COUNT and NUM_NODES and the echo
statements referencing BEST_VAR_NAME, BEST_NODES_STR, SLURM_JOB_NODELIST,
SLURM_NODELIST, and SLURM_STEP_NODELIST, and remove the extra leading whitespace
on the lines echoing SLURM_* and BEST_NODES_STR so formatting is consistent.
- Around line 966-972: The modulo expression ATTEMPT % (60 /
HEALTH_CHECK_INTERVAL) can divide by zero when HEALTH_CHECK_INTERVAL >= 60; fix
by computing a safe divisor first (e.g. LOG_EVERY = 60 / HEALTH_CHECK_INTERVAL
and if LOG_EVERY < 1 set LOG_EVERY=1) and then use ATTEMPT % LOG_EVERY for the
logging condition; update the block using the variables ATTEMPT,
HEALTH_CHECK_INTERVAL, LOG_EVERY (or similar) and keep the existing ELAPSED and
echo logic unchanged.
- Around line 565-572: The health-check loop uses an unquoted command
substitution which can break if curl returns empty or contains whitespace; fix
the loop in the VLLM server wait logic by capturing the curl result into a
quoted variable and comparing it safely, e.g. status="$(curl -s -o /dev/null -w
"%{http_code}" "$VLLM_SERVER_URL/models" 2>/dev/null)" and then use a quoted
comparison like while [ "$status" != "200" ]; do sleep 10; status="$(...)" done;
ensure the variable and substitutions (VLLM_SERVER_URL and the curl command) are
wrapped in double quotes to be defensive.
🧹 Nitpick comments (5)
nemo_skills/pipeline/utils/packager.py (1)
139-157: Mutually reinforcing flags — consider documenting interaction.When
force_pattern_packager=True,repo_pathisNoneso we skip theif repo_path:branch entirely (going to theelseat line 164). Whenforce_installed_nemo_skills=Truebutforce_pattern_packager=False, we enter the git repo branch but include the full installed package. These two flags interact subtly — both are documented in the comments above, which is sufficient, but worth noting thatforce_installed_nemo_skillshas no effect whenforce_pattern_packageris also set (since the code path that checks it is skipped).nemo_skills/pipeline/utils/declarative.py (1)
263-294: PYTHONPATH filtering only removes exact/nemo_run/code— intentional?Line 280 uses
grep -v '^/nemo_run/code$'which only removes the exact path/nemo_run/codefrom PYTHONPATH. If PYTHONPATH contains/nemo_run/code/(trailing slash) or/nemo_run/code/subdir, those entries would be preserved. This is likely intentional (precise removal), but worth confirming it matches the expected behavior. Subdirectory entries like/nemo_run/code/nemo_skillswould also remain.nemo_skills/pipeline/nemo_gym_rollouts.py (2)
184-206: Validation logic is thorough and covers key mutual-exclusion cases.The self-hosted vs. pre-hosted validation gates are well-structured. One edge case: if a user provides
--server_typewith--server_address(pre-hosted), theserver_typeis silently unused. Consider warning or erroring when--server_typeis provided alongside--server_address.
247-255: Broad exception catch is acceptable here but could be narrower.The
except Exceptioncatches all failures when checking remote file existence. Since this involves SSH/cluster operations that can fail unpredictably, the defensive approach with a warning log and safe fallback (include the seed) is pragmatic. If you want to tighten it, catching(OSError, RuntimeError)would cover network/SSH failures while still letting truly unexpected errors surface.nemo_skills/pipeline/utils/scripts.py (1)
690-691: Remove unusedcluster_configfield or document its intended purpose.The
cluster_config: Dictparameter is declared but never used anywhere inMultiVLLMServerScript—not in__post_init__or elsewhere in the class. UnlikeServerScriptandSandboxScript, which both pass this to their respective command builders,MultiVLLMServerScriptbuilds its bash script using only other fields (model_path, num_nodes, gpus_per_server, etc.). Per coding guidelines, avoid silently ignoring user-passed parameters; either remove the field or add clear documentation if it's reserved for future use.
| # Note: --kill-on-bad-exit in srun ensures job fails if vLLM crashes | ||
| VLLM_SERVER_URL="{vllm_server_url}" | ||
| if [ -n "$VLLM_SERVER_URL" ]; then | ||
| echo "=== Waiting for vLLM server at $VLLM_SERVER_URL ===" | ||
| while [ $(curl -s -o /dev/null -w "%{{http_code}}" "$VLLM_SERVER_URL/models" 2>/dev/null) -ne 200 ]; do | ||
| sleep 10 | ||
| done | ||
| echo "vLLM server is ready!" |
There was a problem hiding this comment.
Unquoted command substitution in health-check loop.
Line 569: The $(curl ...) is not double-quoted. While curl -w "%{http_code}" normally returns "000" on connection failure, quoting defensively is safer.
🛡️ Suggested fix
- while [ $(curl -s -o /dev/null -w "%{{http_code}}" "$VLLM_SERVER_URL/models" 2>/dev/null) -ne 200 ]; do
+ while [ "$(curl -s -o /dev/null -w "%{{http_code}}" "$VLLM_SERVER_URL/models" 2>/dev/null)" -ne 200 ]; do🤖 Prompt for AI Agents
In `@nemo_skills/pipeline/utils/scripts.py` around lines 565 - 572, The
health-check loop uses an unquoted command substitution which can break if curl
returns empty or contains whitespace; fix the loop in the VLLM server wait logic
by capturing the curl result into a quoted variable and comparing it safely,
e.g. status="$(curl -s -o /dev/null -w "%{http_code}" "$VLLM_SERVER_URL/models"
2>/dev/null)" and then use a quoted comparison like while [ "$status" != "200"
]; do sleep 10; status="$(...)" done; ensure the variable and substitutions
(VLLM_SERVER_URL and the curl command) are wrapped in double quotes to be
defensive.
| fi | ||
| # Log progress every 60 seconds | ||
| if [ $((ATTEMPT % (60 / HEALTH_CHECK_INTERVAL))) -eq 0 ]; then | ||
| ELAPSED=$(($(date +%s) - START_TIME)) | ||
| echo " ... still waiting for server $GLOBAL_IDX (${{ELAPSED}}s elapsed)" | ||
| fi | ||
| sleep $HEALTH_CHECK_INTERVAL |
There was a problem hiding this comment.
Division-by-zero risk if health_check_interval ≥ 60.
Line 968: ATTEMPT % (60 / HEALTH_CHECK_INTERVAL) — bash integer division truncates, so if HEALTH_CHECK_INTERVAL > 60, 60 / HEALTH_CHECK_INTERVAL evaluates to 0, causing % 0 which is a bash arithmetic error. The default of 10 is safe, but this is a latent issue if someone configures a large interval.
🛡️ Suggested fix
- if [ $((ATTEMPT % (60 / HEALTH_CHECK_INTERVAL))) -eq 0 ]; then
+ LOG_EVERY=$((60 / HEALTH_CHECK_INTERVAL))
+ if [ "$LOG_EVERY" -le 0 ]; then LOG_EVERY=1; fi
+ if [ $((ATTEMPT % LOG_EVERY)) -eq 0 ]; then📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| fi | |
| # Log progress every 60 seconds | |
| if [ $((ATTEMPT % (60 / HEALTH_CHECK_INTERVAL))) -eq 0 ]; then | |
| ELAPSED=$(($(date +%s) - START_TIME)) | |
| echo " ... still waiting for server $GLOBAL_IDX (${{ELAPSED}}s elapsed)" | |
| fi | |
| sleep $HEALTH_CHECK_INTERVAL | |
| fi | |
| # Log progress every 60 seconds | |
| LOG_EVERY=$((60 / HEALTH_CHECK_INTERVAL)) | |
| if [ "$LOG_EVERY" -le 0 ]; then LOG_EVERY=1; fi | |
| if [ $((ATTEMPT % LOG_EVERY)) -eq 0 ]; then | |
| ELAPSED=$(($(date +%s) - START_TIME)) | |
| echo " ... still waiting for server $GLOBAL_IDX (${{ELAPSED}}s elapsed)" | |
| fi | |
| sleep $HEALTH_CHECK_INTERVAL |
🤖 Prompt for AI Agents
In `@nemo_skills/pipeline/utils/scripts.py` around lines 966 - 972, The modulo
expression ATTEMPT % (60 / HEALTH_CHECK_INTERVAL) can divide by zero when
HEALTH_CHECK_INTERVAL >= 60; fix by computing a safe divisor first (e.g.
LOG_EVERY = 60 / HEALTH_CHECK_INTERVAL and if LOG_EVERY < 1 set LOG_EVERY=1) and
then use ATTEMPT % LOG_EVERY for the logging condition; update the block using
the variables ATTEMPT, HEALTH_CHECK_INTERVAL, LOG_EVERY (or similar) and keep
the existing ELAPSED and echo logic unchanged.
- Fix PYTHONPATH stripping to match /nemo_run/code prefixed paths - Re-enable pipefail before ng_collect_rollouts - Fix indentation in GymClientScript error block - Remove emoji from server ready message Signed-off-by: George Armstrong <georgea@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
nemo_skills/pipeline/utils/declarative.py (1)
778-802:⚠️ Potential issue | 🟠 MajorBug:
external_depsis not passed to the first executor when reordering occurs.After the sort at line 754,
prepared_commandsis reordered so spanning components appear first (improving SLURM allocation). However, the loop at line 762 checkshet_idx == 0 and comp_idx == 0using the original component indices stored in each entry. If the spanning component that now appears first was originally atcomp_idx=1, the condition fails and no executor receivesexternal_deps.The code already acknowledges this reordering issue at lines 757–759 for
shared_packager, but the same problem applies toexternal_deps. Since nemo-run uses the first executor to determine the SLURM allocation (line 738), the fix is to pass dependencies to the first executor in iteration order:- for entry in prepared_commands: + for entry_idx, entry in enumerate(prepared_commands): het_idx = entry["het_idx"] comp_idx = entry["comp_idx"] ... - exec_dependencies = external_deps if (het_idx == 0 and comp_idx == 0) else None + exec_dependencies = external_deps if entry_idx == 0 else None🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nemo_skills/pipeline/utils/declarative.py` around lines 778 - 802, prepared_commands is reordered so relying on stored het_idx/comp_idx to decide who gets external_deps is incorrect; instead, determine the "first executor" by iteration order and pass external_deps to that executor. Modify the loop that builds executors (the block that calls _create_executor) to compute exec_dependencies based on the iteration position (e.g., using an iteration counter or a boolean flag) rather than the stored het_idx/comp_idx, ensuring the very first created executor receives external_deps; update any existing shared_packager handling consistently so only the first yielded executor gets external_deps.
🧹 Nitpick comments (4)
nemo_skills/pipeline/utils/declarative.py (1)
638-644: Use direct dictionary access forexec_config["mounts"].
exec_config["mounts"]is always set at line 302, so.get()is unnecessary. Per coding guidelines, use direct access to fail fast if the key is ever unexpectedly missing.Suggested fix
- extra_mounts = exec_config.get("mounts") or None + extra_mounts = exec_config["mounts"] or NoneAs per coding guidelines, "Do not use
.get()for accessing dictionary keys if the code expects them to be present; use direct dictionary accessdict[key]instead to allow proper error handling and fail fast with clear errors".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nemo_skills/pipeline/utils/declarative.py` around lines 638 - 644, Replace the use of exec_config.get("mounts") with direct dictionary access exec_config["mounts"] in the mounts merge block so the code fails fast if the key is missing; specifically, in the block that computes mounts (using variables mounts, extra_mounts and calling get_mounts_from_config(cluster_config)), set extra_mounts = exec_config["mounts"] (keeping the subsequent truthy check and merge logic unchanged) to follow the guideline of direct access when the key is expected to exist.nemo_skills/pipeline/utils/scripts.py (3)
518-535:extra_argumentsis forwarded verbatim to bothng_runandng_collect_rollouts.If a user passes Hydra overrides that are only valid for one of the two commands, the other will fail with an unknown-override error. Consider splitting into separate parameters (e.g.,
extra_ng_run_args/extra_ng_collect_args) or at least document the coupling more prominently so callers know every override must be understood by both commands.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nemo_skills/pipeline/utils/scripts.py` around lines 518 - 535, extra_arguments is forwarded verbatim to both ng_run and ng_collect_rollouts causing unknown-override failures; split or separate the args so each command only receives relevant overrides. Change the implementation to accept two distinct parameters (e.g., extra_ng_run_args and extra_ng_collect_args) and use extra_ng_run_args when appending to ng_run_parts and building ng_run_cmd, and use extra_ng_collect_args when appending to ng_collect_parts and building ng_collect_cmd; update constructor/attribute names and any call sites that set extra_arguments to pass the appropriate new parameter(s), or alternatively update documentation to clearly state that extra_arguments must be valid for both commands if you choose not to split.
639-805:MultiVLLMServerScript— solid design for data-parallel vLLM.Clean use of SLURM env vars for per-task GPU slicing and port assignment. A few observations:
--trust-remote-codeis unconditionally hardcoded (line 800). If a user loads a model from a trusted local checkpoint and wants to avoid this flag, there's no way to opt out. Consider making it configurable or letting users control it viaserver_args.- The port range (
base_porttobase_port + total_replicas - 1) is allocated once at init time. If any of those ports are already in use on a target node, the server will fail at runtime. This is inherent to the approach but worth noting in the docstring.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nemo_skills/pipeline/utils/scripts.py` around lines 639 - 805, The script unconditionally injects --trust-remote-code in MultiVLLMServerScript which prevents users from opting out; add a boolean field (e.g., trust_remote_code: bool = True) to the dataclass and conditionally include the "--trust-remote-code" token in the generated cmd only when that flag is true (or let users pass it via server_args by default). Update the __post_init__ build logic that constructs cmd to reference this new field when composing the python3 -m vllm.entrypoints.openai.api_server command, and update the class docstring to note the port allocation caveat (base_port .. base_port + _total_replicas - 1) so users know ports may be in-use on target nodes.
625-633:hostname_ref()in env vars won't resolve in heterogeneous jobs.Line 630 sets
NEMO_SKILLS_SANDBOX_HOSTtoself.sandbox.hostname_ref(). For non-het jobs this returns"127.0.0.1"(fine), but for het jobs it returns a bash expression like${SLURM_MASTER_NODE_HET_GROUP_0:-localhost}. Since env vars set via nemo-run are literal values (not shell-evaluated), OmegaConf/Gym would receive the unexpanded string.This isn't an issue for the current single-group use case, but will silently break if someone uses this script in a heterogeneous job. Consider adding a guard or a comment.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nemo_skills/pipeline/utils/scripts.py` around lines 625 - 633, The code sets NEMO_SKILLS_SANDBOX_HOST to self.sandbox.hostname_ref() which can be a bash expression for heterogeneous (het) SLURM jobs and will be passed literally; change the logic around env_vars population so you detect bash-expression patterns (e.g. contains "${" or starts with "$") from self.sandbox.hostname_ref() and avoid injecting an unexpanded literal into env_vars["NEMO_SKILLS_SANDBOX_HOST"] — instead either (a) set a safe fallback like "localhost"/resolved IP, (b) omit the env var so the target process can resolve itself, and (c) add a warning log entry; keep env_vars["NEMO_SKILLS_SANDBOX_PORT"] as-is. Update the code around self.sandbox.hostname_ref(), self.sandbox.port and the return that builds {"environment": env_vars} and add a brief comment noting het-only hostname expressions must not be passed literally.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@nemo_skills/pipeline/utils/scripts.py`:
- Around line 590-605: The health check currently relies on a brittle string
match of STATUS_OUTPUT for "healthy, 0 unhealthy" (in the ng_status loop) which
can break if ng_status output changes; update the loop to use a more robust
signal: prefer ng_status's exit code or a structured output (JSON) if available,
or parse numeric healthy/unhealthy counts with a regex that extracts both
numbers from STATUS_OUTPUT and checks unhealthy == 0; also add a timeout/counter
to the loop to prevent an infinite hang, referencing the ng_status call,
STATUS_OUTPUT, CURRENT_STATUS and LAST_STATUS variables for where to implement
the changes.
---
Outside diff comments:
In `@nemo_skills/pipeline/utils/declarative.py`:
- Around line 778-802: prepared_commands is reordered so relying on stored
het_idx/comp_idx to decide who gets external_deps is incorrect; instead,
determine the "first executor" by iteration order and pass external_deps to that
executor. Modify the loop that builds executors (the block that calls
_create_executor) to compute exec_dependencies based on the iteration position
(e.g., using an iteration counter or a boolean flag) rather than the stored
het_idx/comp_idx, ensuring the very first created executor receives
external_deps; update any existing shared_packager handling consistently so only
the first yielded executor gets external_deps.
---
Duplicate comments:
In `@nemo_skills/pipeline/utils/scripts.py`:
- Around line 565-572: The health-check loop uses unquoted command substitution
for the curl exit code; update the while condition in the VLLM_SERVER_URL block
so the $(curl ...) expansion is wrapped in double quotes to prevent
word-splitting/empty-string issues (i.e., quote the command substitution used in
the while test that queries "$VLLM_SERVER_URL/models" so the comparison against
200 is reliable).
- Around line 933-939: The error block around the BEST_COUNT check has
consistent 4-space indentation for the echo lines but the review contains a
duplicate comment marker; remove the duplicate marker and verify the conditional
block (if [ "$BEST_COUNT" -lt "$NUM_NODES" ]; then ... exit 1) keeps the echo
lines for BEST_VAR_NAME, BEST_NODES_STR, SLURM_JOB_NODELIST, SLURM_NODELIST, and
SLURM_STEP_NODELIST indented consistently and that no stray duplicate comments
remain in scripts.py.
- Around line 963-976: The modulo expression ATTEMPT % (60 /
HEALTH_CHECK_INTERVAL) can divide by zero when HEALTH_CHECK_INTERVAL >= 60; fix
by computing a safe log-frequency before the loop (e.g. LOG_EVERY=$((60 /
HEALTH_CHECK_INTERVAL)); if [ "$LOG_EVERY" -le 0 ]; then LOG_EVERY=1; fi) and
then replace the existing check with if [ $((ATTEMPT % LOG_EVERY)) -eq 0 ]; then
...; use the same variables ATTEMPT, HEALTH_CHECK_INTERVAL, LOG_EVERY, ELAPSED
and GLOBAL_IDX to locate and update the code inside the while true loop.
---
Nitpick comments:
In `@nemo_skills/pipeline/utils/declarative.py`:
- Around line 638-644: Replace the use of exec_config.get("mounts") with direct
dictionary access exec_config["mounts"] in the mounts merge block so the code
fails fast if the key is missing; specifically, in the block that computes
mounts (using variables mounts, extra_mounts and calling
get_mounts_from_config(cluster_config)), set extra_mounts =
exec_config["mounts"] (keeping the subsequent truthy check and merge logic
unchanged) to follow the guideline of direct access when the key is expected to
exist.
In `@nemo_skills/pipeline/utils/scripts.py`:
- Around line 518-535: extra_arguments is forwarded verbatim to both ng_run and
ng_collect_rollouts causing unknown-override failures; split or separate the
args so each command only receives relevant overrides. Change the implementation
to accept two distinct parameters (e.g., extra_ng_run_args and
extra_ng_collect_args) and use extra_ng_run_args when appending to ng_run_parts
and building ng_run_cmd, and use extra_ng_collect_args when appending to
ng_collect_parts and building ng_collect_cmd; update constructor/attribute names
and any call sites that set extra_arguments to pass the appropriate new
parameter(s), or alternatively update documentation to clearly state that
extra_arguments must be valid for both commands if you choose not to split.
- Around line 639-805: The script unconditionally injects --trust-remote-code in
MultiVLLMServerScript which prevents users from opting out; add a boolean field
(e.g., trust_remote_code: bool = True) to the dataclass and conditionally
include the "--trust-remote-code" token in the generated cmd only when that flag
is true (or let users pass it via server_args by default). Update the
__post_init__ build logic that constructs cmd to reference this new field when
composing the python3 -m vllm.entrypoints.openai.api_server command, and update
the class docstring to note the port allocation caveat (base_port .. base_port +
_total_replicas - 1) so users know ports may be in-use on target nodes.
- Around line 625-633: The code sets NEMO_SKILLS_SANDBOX_HOST to
self.sandbox.hostname_ref() which can be a bash expression for heterogeneous
(het) SLURM jobs and will be passed literally; change the logic around env_vars
population so you detect bash-expression patterns (e.g. contains "${" or starts
with "$") from self.sandbox.hostname_ref() and avoid injecting an unexpanded
literal into env_vars["NEMO_SKILLS_SANDBOX_HOST"] — instead either (a) set a
safe fallback like "localhost"/resolved IP, (b) omit the env var so the target
process can resolve itself, and (c) add a warning log entry; keep
env_vars["NEMO_SKILLS_SANDBOX_PORT"] as-is. Update the code around
self.sandbox.hostname_ref(), self.sandbox.port and the return that builds
{"environment": env_vars} and add a brief comment noting het-only hostname
expressions must not be passed literally.
| STATUS_OUTPUT=$(ng_status 2>&1) | ||
|
|
||
| if echo "$STATUS_OUTPUT" | grep -q "healthy, 0 unhealthy"; then | ||
| echo "All servers ready!" | ||
| break | ||
| fi | ||
|
|
||
| # Only print status when it changes (reduce verbosity) | ||
| CURRENT_STATUS=$(echo "$STATUS_OUTPUT" | grep -oE '[0-9]+ healthy' | head -1 || echo "starting") | ||
| if [ "$CURRENT_STATUS" != "$LAST_STATUS" ]; then | ||
| echo "Server status: $CURRENT_STATUS" | ||
| LAST_STATUS="$CURRENT_STATUS" | ||
| fi | ||
|
|
||
| sleep 10 | ||
| done |
There was a problem hiding this comment.
Fragile health-check string match: "healthy, 0 unhealthy".
If ng_status changes its output format (e.g., wording, ordering, locale), the grep -q "healthy, 0 unhealthy" at line 592 will never match and the loop will spin forever (silent hang with no timeout). Consider using a more structured health-check mechanism (e.g., exit code of ng_status, or a JSON output field) if one is available.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@nemo_skills/pipeline/utils/scripts.py` around lines 590 - 605, The health
check currently relies on a brittle string match of STATUS_OUTPUT for "healthy,
0 unhealthy" (in the ng_status loop) which can break if ng_status output
changes; update the loop to use a more robust signal: prefer ng_status's exit
code or a structured output (JSON) if available, or parse numeric
healthy/unhealthy counts with a regex that extracts both numbers from
STATUS_OUTPUT and checks unhealthy == 0; also add a timeout/counter to the loop
to prevent an infinite hang, referencing the ng_status call, STATUS_OUTPUT,
CURRENT_STATUS and LAST_STATUS variables for where to implement the changes.
|
curious if this will support resuming rollouts like with .done file? |
@cmunley1 no concrete plans at the moment, it would be better if Gym supported resuming natively. |
|
im also curious if this would/could enable an end to end reward profiling across many parallel slurm jobs ie say I want to run 16 rollouts per prompt, for 20,000 prompts, across a few hundred slurm jobs, then aggregate results, so i have avg/stddev reward per task for filtering or curriculum we have a small feature in nemo gym to aggregate results, but chunking across slurm jobs is not supported in gym natively, it feels more suited here since ns can launch slurm jobs This is a common workflow that currently requires custom scripts in nemo gym to chunk, submit parallel jobs, and merge result |
|
@cmunley1 we do something similar to |
Kipok
left a comment
Kipok
left a comment
There was a problem hiding this comment.
added a few comments. Here is another one from codex, not fully sure if it's meaningful or not, but worth checking
Sorting prepared_commands changes executor order, but external dependencies are still attached based on the original comp_idx == 0 check later in the loop. When a non-spanning command was originally first and a spanning command is reordered ahead of it, the dependency can be attached to a non-leading executor, which can let a run_after job start before its intended upstream dependency.
|
|
||
|
|
||
| @dataclass(kw_only=True) | ||
| class NemoGymRolloutsScript(BaseJobScript): |
There was a problem hiding this comment.
might be a good time to make scripts a subfolder and split this into different files
There was a problem hiding this comment.
fair, let's do that
|
|
||
|
|
||
| @dataclass(kw_only=True) | ||
| class MultiVLLMServerScript(BaseJobScript): |
There was a problem hiding this comment.
I'm a bit confused why we need this and below class. vllm should naturally support data parallel with a --data-parallel-size argument. Why can't we just use that with our normal vllm server? And then there is a single ip gym talks to, but it still handled data split, just on the server side?
There was a problem hiding this comment.
Yeah that was definitely the preferred route, but I personally found the data parallel size to be error prone and not work very well past a couple nodes. I thought it might be useful to share, but this is unused in the main pipeline at the moment, so I should probably just remove it.
| cmd = f"""set -e | ||
| set -o pipefail | ||
|
|
||
| echo "=== Installing NeMo Gym ===" |
There was a problem hiding this comment.
why do we need to install gym? It should be preinstalled in the container, right?
There was a problem hiding this comment.
This is needed when users mount a custom Gym path (e.g., dev branch) over the container's built-in path. The mounted directory may not have a .venv. Added a comment explaining this. The --allow-existing flag makes it fast (~1s) when the venv is already up to date.
| # Wait for vLLM server to be ready before starting ng_run | ||
| # Note: --kill-on-bad-exit in srun ensures job fails if vLLM crashes | ||
| VLLM_SERVER_URL="{vllm_server_url}" | ||
| if [ -n "$VLLM_SERVER_URL" ]; then |
There was a problem hiding this comment.
can we reuse get_server_wait_cmd here?
There was a problem hiding this comment.
Done in 402715f. Replaced with get_server_wait_cmd() from nemo_skills.utils.
| "dummy", | ||
| help="API key for policy server. Use 'dummy' for local vLLM servers.", | ||
| ), | ||
| policy_model_name: str = typer.Option( |
There was a problem hiding this comment.
can we just reuse the "model" parameter here? That would be consistent with how we do that in main generate pipeline
There was a problem hiding this comment.
I'm not sure if what you're saying is completely consistent with ns generate. But I aligned the behavior here with ns generate in 12169c5. --model is the model path/name, --server_address is the pre-hosted URL, and --server_gpus is the discriminator (truthy = self-hosted, falsy = pre-hosted). This matches configure_client() in generation.py where if server_gpus: determines hosting mode.
| if random_seeds is not None: | ||
| # Explicit seeds provided | ||
| if isinstance(random_seeds, str): | ||
| seed_indices = [int(s.strip()) for s in random_seeds.split(",")] |
There was a problem hiding this comment.
can we reuse str_ids_to_list from main generate pipeline for consistency?
There was a problem hiding this comment.
Done in 402715f. Using str_ids_to_list() now.
| skipped_seeds.append(seed_idx) | ||
| else: | ||
| filtered_seeds.append(seed_idx) | ||
| except Exception as e: |
There was a problem hiding this comment.
do we expect this to fail? Why blank exception?
There was a problem hiding this comment.
Done in 402715f. Removed the try/except — cluster check failures now propagate immediately.
- Split scripts.py into scripts/ package (base, server, generation, nemo_gym) - Remove unused MultiVLLMServerScript and GymClientScript classes - Remove server_address param; model now accepts URLs for pre-hosted servers - Use str_ids_to_list for seed parsing (consistency with generate pipeline) - Add server_container CLI option; use direct dict access for container lookup - Remove try/except on seed skip check — fail hard on cluster errors - Use get_server_wait_cmd for vLLM health check (reuse existing utility) - Add comment explaining why Gym venv install step is needed for mounted paths - Fix external_deps assignment after prepared_commands reordering - Replace .get() with direct dict access throughout Signed-off-by: George Armstrong <georgea@nvidia.com>
Signed-off-by: George Armstrong <georgea@nvidia.com>
Signed-off-by: George Armstrong <georgea@nvidia.com>
Keep --server_address as separate param (matching generate pipeline). Use server_gpus as the discriminator: truthy = self-hosted, falsy = pre-hosted. This matches configure_client() in generation.py. Signed-off-by: George Armstrong <georgea@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (1)
nemo_skills/pipeline/utils/scripts/server.py (1)
136-142: Use direct key access for required metadata fields.
metadata["environment"]is expected by contract here; using.get()can silently hide shape regressions.As per coding guidelines: "Don't use `.get()` for accessing dictionary keys if the code expects them to be present; use direct access `data[key_name]` to fail with a clear error instead of silently corrupting data".♻️ Proposed fix
def build_cmd() -> Tuple[str, Dict]: - env = dict(metadata.get("environment", {})) + env = dict(metadata["environment"]) if self.env_overrides: for override in self.env_overrides: key, value = override.split("=", 1) env[key] = value🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nemo_skills/pipeline/utils/scripts/server.py` around lines 136 - 142, The code in build_cmd uses metadata.get("environment", {}) which can hide missing required fields; change this to directly access metadata["environment"] (e.g., env = dict(metadata["environment"])) so a KeyError is raised if the environment key is absent, then apply self.env_overrides as before and return cmd, {"environment": env}; reference build_cmd, metadata, and self.env_overrides when making the change.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@nemo_skills/pipeline/nemo_gym_rollouts.py`:
- Around line 188-202: When model is a pre-hosted URL (pre_hosted == True),
reject any self-hosted-only flags instead of silently ignoring them: detect if
any of the parameters server_type, server_gpus, server_nodes, server_args, or
server_container are non-None/used while pre_hosted is True and raise a
ValueError stating those flags are invalid for pre-hosted mode; update the
initial validation block around pre_hosted/self_hosted (and the similar
validation later in the same file where server_* flags are checked) to perform
this check and include the offending flag names in the error message so users
fail fast on misconfiguration.
- Around line 221-233: The seed-selection logic around variables random_seeds,
num_random_seeds, seed_indices and helper str_ids_to_list needs validation:
explicitly validate num_random_seeds is an int > 0 (reject or raise/log error
for 0 or negative values) instead of relying on truthiness, and when parsing
random_seeds ensure you convert entries to ints, deduplicate them (preserve
order) to avoid duplicate job/output names, and reject/normalize
invalid/non-negative seeds; apply the same validation/deduplication fix to the
other occurrences mentioned (the blocks using seed_indices at lines referred to)
so all job-generation code uses the validated, deduplicated list or errors out
with a clear message.
In `@nemo_skills/pipeline/utils/scripts/generation.py`:
- Around line 86-94: The current construction of server_addresses can ignore or
misindex pre-hosted addresses because server_addresses_prehosted is only used
when self.servers is present and there are no length checks; update the logic
that builds server_addresses (the block using self.servers,
server_script.hostname_ref(), server_script.port,
self.server_addresses_prehosted and enumerate) to validate inputs: enforce that
self.servers and self.server_addresses_prehosted are mutually exclusive or
explicitly allowed, check that when both are intended to be used their lengths
match, and guard the access to self.server_addresses_prehosted[server_idx] with
an index check — otherwise raise a clear ValueError indicating the mismatch or
unsupported combination so caller errors rather than producing opaque failures.
In `@nemo_skills/pipeline/utils/scripts/nemo_gym.py`:
- Around line 62-80: The __post_init__ currently silently ignores server_address
when server is also set; update the initialization to validate inputs and fail
fast: in __post_init__ (before build_cmd is used) check if both server and
server_address are provided and raise a ValueError (or TypeError) with a clear
message, referencing the conflicting fields server and server_address so callers
know to pass only one; keep the existing build_cmd logic unchanged but rely on
this early validation to prevent silent shadowing when building ng_run_parts and
appending '+policy_base_url=...'.
In `@nemo_skills/pipeline/utils/scripts/server.py`:
- Around line 126-134: The __post_init__ currently allows calling
sandbox_command with port=None when allocate_port is False, causing invalid
runtime env and delayed failures; update __post_init__ (in the class that
defines allocate_port and port) to validate that if allocate_port is False then
port is not None and raise a clear ValueError (or similar) immediately,
otherwise if allocate_port is True keep allocating via
get_free_port(strategy="random") and proceed to call
sandbox_command(cluster_config=self.cluster_config, port=self.port); this
ensures failures are fast and surfaces the missing-port configuration early.
- Around line 79-93: In __post_init__, add a fail-fast check that raises a clear
exception if allocate_port is False and port is None: if not self.allocate_port
and self.port is None raise a ValueError (or custom exception) with a message
like "port must be provided when allocate_port is False" before calling
get_server_command; this ensures callers of server.__post_init__ (and the
get_server_command call) don't proceed with a None server_port.
---
Nitpick comments:
In `@nemo_skills/pipeline/utils/scripts/server.py`:
- Around line 136-142: The code in build_cmd uses metadata.get("environment",
{}) which can hide missing required fields; change this to directly access
metadata["environment"] (e.g., env = dict(metadata["environment"])) so a
KeyError is raised if the environment key is absent, then apply
self.env_overrides as before and return cmd, {"environment": env}; reference
build_cmd, metadata, and self.env_overrides when making the change.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (11)
docs/basics/code-packaging.mdnemo_skills/pipeline/nemo_gym_rollouts.pynemo_skills/pipeline/utils/__init__.pynemo_skills/pipeline/utils/declarative.pynemo_skills/pipeline/utils/scripts.pynemo_skills/pipeline/utils/scripts/__init__.pynemo_skills/pipeline/utils/scripts/base.pynemo_skills/pipeline/utils/scripts/generation.pynemo_skills/pipeline/utils/scripts/nemo_gym.pynemo_skills/pipeline/utils/scripts/server.pytests/gpu-tests/test_nemo_gym_rollouts.py
💤 Files with no reviewable changes (1)
- nemo_skills/pipeline/utils/scripts.py
| server_addresses = None | ||
| if self.servers is not None: | ||
| server_addresses = [] | ||
| for server_idx, server_script in enumerate(self.servers): | ||
| if server_script is not None: | ||
| addr = f"{server_script.hostname_ref()}:{server_script.port}" | ||
| else: | ||
| addr = self.server_addresses_prehosted[server_idx] | ||
| server_addresses.append(addr) |
There was a problem hiding this comment.
Pre-hosted server addresses can be ignored or misindexed.
Right now, pre-hosted addresses are only used when servers is provided, and missing length checks can produce opaque failures.
🔧 Proposed fix
server_addresses = None
if self.servers is not None:
+ if any(server is None for server in self.servers):
+ if self.server_addresses_prehosted is None:
+ raise ValueError(
+ "server_addresses_prehosted is required when `servers` contains None entries."
+ )
+ if len(self.server_addresses_prehosted) != len(self.servers):
+ raise ValueError("server_addresses_prehosted must match the length of `servers`.")
server_addresses = []
for server_idx, server_script in enumerate(self.servers):
if server_script is not None:
addr = f"{server_script.hostname_ref()}:{server_script.port}"
else:
addr = self.server_addresses_prehosted[server_idx]
server_addresses.append(addr)
+ elif self.server_addresses_prehosted is not None:
+ server_addresses = list(self.server_addresses_prehosted)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@nemo_skills/pipeline/utils/scripts/generation.py` around lines 86 - 94, The
current construction of server_addresses can ignore or misindex pre-hosted
addresses because server_addresses_prehosted is only used when self.servers is
present and there are no length checks; update the logic that builds
server_addresses (the block using self.servers, server_script.hostname_ref(),
server_script.port, self.server_addresses_prehosted and enumerate) to validate
inputs: enforce that self.servers and self.server_addresses_prehosted are
mutually exclusive or explicitly allowed, check that when both are intended to
be used their lengths match, and guard the access to
self.server_addresses_prehosted[server_idx] with an index check — otherwise
raise a clear ValueError indicating the mismatch or unsupported combination so
caller errors rather than producing opaque failures.
Good catch. Done in |
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
nemo_skills/pipeline/nemo_gym_rollouts.py (2)
221-233:⚠️ Potential issue | 🟠 MajorValidate seed inputs explicitly to avoid empty runs and output collisions.
elif num_random_seeds:relies on truthiness (0falls through; negatives can create empty ranges), and explicit seed lists are not validated for duplicates/non-negative values. This can lead to duplicate job/output names downstream (rollouts-rs{seed}.jsonl) or silently creating zero jobs.Suggested fix
if random_seeds is not None: if isinstance(random_seeds, str): seed_indices = str_ids_to_list(random_seeds) else: seed_indices = list(random_seeds) + if not seed_indices: + raise ValueError("--random_seeds must contain at least one seed") + if any(seed < 0 for seed in seed_indices): + raise ValueError("--random_seeds must contain non-negative integers") + if len(seed_indices) != len(set(seed_indices)): + raise ValueError("--random_seeds must not contain duplicate values") LOG.info(f"Using explicit seeds: {seed_indices}") - elif num_random_seeds: + elif num_random_seeds is not None: + if num_random_seeds <= 0: + raise ValueError("--num_random_seeds must be > 0") + if starting_seed < 0: + raise ValueError("--starting_seed must be >= 0") seed_indices = list(range(starting_seed, starting_seed + num_random_seeds)) LOG.info( f"Creating {num_random_seeds} separate jobs (rs{starting_seed}..rs{starting_seed + num_random_seeds - 1})" )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nemo_skills/pipeline/nemo_gym_rollouts.py` around lines 221 - 233, The seed-selection logic around random_seeds/num_random_seeds should validate inputs: ensure num_random_seeds is an int >= 1 (do not rely on truthiness) and starting_seed is a non-negative int before constructing seed_indices, and when random_seeds is provided (handled by str_ids_to_list/random_seeds branch) convert elements to ints, check they are non-negative and deduplicate (or raise ValueError) to avoid duplicate job names and empty ranges; update the block that sets seed_indices (using random_seeds, str_ids_to_list, num_random_seeds, starting_seed) to perform these checks and raise/return a clear error if validation fails instead of silently producing empty or duplicate seed lists.
189-204:⚠️ Potential issue | 🟠 MajorReject mode-incompatible flags instead of silently ignoring them.
In pre-hosted mode, self-hosted-only flags (
--server_type,--server_gpus,--server_nodes,--server_args,--server_container) are currently accepted but ignored; similarly,--server_addressshould be rejected in self-hosted mode. Fail fast here to prevent silent misconfiguration.Suggested fix
if pre_hosted and server_address is None: raise ValueError("--server_address is required when not using self-hosted server (no --server_gpus)") + + if pre_hosted: + invalid_flags = [] + if server_type is not None: + invalid_flags.append("--server_type") + if server_gpus is not None: + invalid_flags.append("--server_gpus") + if server_nodes != 1: + invalid_flags.append("--server_nodes") + if server_args: + invalid_flags.append("--server_args") + if server_container is not None: + invalid_flags.append("--server_container") + if invalid_flags: + raise ValueError(f"{', '.join(invalid_flags)} are only valid in self-hosted mode") + + if self_hosted and server_address is not None: + raise ValueError("--server_address cannot be used with self-hosted mode")As per coding guidelines: "Avoid cases where user-passed parameters are unused; code should fail if user specifies an unsupported argument or if a required argument is missing. Use dataclass or **kwargs syntax to handle this automatically".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nemo_skills/pipeline/nemo_gym_rollouts.py` around lines 189 - 204, The current hosting-mode validation accepts mode-incompatible flags silently; update the checks around self_hosted/pre_hosted (variables self_hosted, pre_hosted) to reject incompatible arguments: when pre_hosted is true, if any of server_type, server_gpus, server_nodes, server_args, or server_container are not None, raise a ValueError naming the offending flags; conversely when self_hosted is true and server_address is not None, raise a ValueError rejecting --server_address; keep the existing required-flag checks for model, server_type and policy_model_name intact but place these new incompatibility checks near the top of the block so the function fails fast on misconfiguration.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@nemo_skills/pipeline/nemo_gym_rollouts.py`:
- Around line 185-187: The parsed config paths list (config_paths_list obtained
from config_paths.split(...)) can be empty after stripping, so add an explicit
fail-fast check right after building config_paths_list: if not
config_paths_list: log an error via LOG.error (or raise a ValueError) with a
clear message like "No valid --config_paths provided; please supply at least one
path" and raise/exit before any job construction or further processing (refer to
the config_paths and config_paths_list variables in nemo_gym_rollouts.py).
---
Duplicate comments:
In `@nemo_skills/pipeline/nemo_gym_rollouts.py`:
- Around line 221-233: The seed-selection logic around
random_seeds/num_random_seeds should validate inputs: ensure num_random_seeds is
an int >= 1 (do not rely on truthiness) and starting_seed is a non-negative int
before constructing seed_indices, and when random_seeds is provided (handled by
str_ids_to_list/random_seeds branch) convert elements to ints, check they are
non-negative and deduplicate (or raise ValueError) to avoid duplicate job names
and empty ranges; update the block that sets seed_indices (using random_seeds,
str_ids_to_list, num_random_seeds, starting_seed) to perform these checks and
raise/return a clear error if validation fails instead of silently producing
empty or duplicate seed lists.
- Around line 189-204: The current hosting-mode validation accepts
mode-incompatible flags silently; update the checks around
self_hosted/pre_hosted (variables self_hosted, pre_hosted) to reject
incompatible arguments: when pre_hosted is true, if any of server_type,
server_gpus, server_nodes, server_args, or server_container are not None, raise
a ValueError naming the offending flags; conversely when self_hosted is true and
server_address is not None, raise a ValueError rejecting --server_address; keep
the existing required-flag checks for model, server_type and policy_model_name
intact but place these new incompatibility checks near the top of the block so
the function fails fast on misconfiguration.
| # Parse config paths | ||
| config_paths_list = [p.strip() for p in config_paths.split(",") if p.strip()] | ||
| LOG.info(f"Config paths: {config_paths_list}") |
There was a problem hiding this comment.
Fail fast when parsed --config_paths is empty.
--config_paths is required, but after parsing/stripping it can still become an empty list (e.g., whitespace/comma-only input). Add an explicit check to raise a clear error before job construction.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@nemo_skills/pipeline/nemo_gym_rollouts.py` around lines 185 - 187, The parsed
config paths list (config_paths_list obtained from config_paths.split(...)) can
be empty after stripping, so add an explicit fail-fast check right after
building config_paths_list: if not config_paths_list: log an error via LOG.error
(or raise a ValueError) with a clear message like "No valid --config_paths
provided; please supply at least one path" and raise/exit before any job
construction or further processing (refer to the config_paths and
config_paths_list variables in nemo_gym_rollouts.py).
- Reject self-hosted flags (server_type, server_args, etc.) in pre-hosted mode - Validate seed inputs: reject num_random_seeds <= 0 and duplicate seeds - Reject conflicting server + server_address in NemoGymRolloutsScript - Fail fast when allocate_port=False and port is None (ServerScript, SandboxScript) Signed-off-by: George Armstrong <georgea@nvidia.com>
Signed-off-by: George Armstrong <georgea@nvidia.com>
gwarmstrong
changed the title
Kipok
left a comment
Kipok
left a comment
There was a problem hiding this comment.
probably worth running slurm tests as well here
3 participants
NeMo-Gym Rollouts Pipeline
Add
ns nemo_gym_rolloutsCLI command for orchestrating rollout collection with NeMo Gym, including self-hosted or pre-hosted vLLM servers, optional sandbox, and parallel seed-based scaling.New script classes:
Infrastructure enhancements:
Gym Config Verification: math_with_judge (AIME24) & code_gen (LCB)
Verified that
nemo_gym_rolloutspipeline works end-to-end withmath_with_judgeandcode_genGym configs on their full validation datasets.Results
Model:
nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16(8x GPU, vLLM,--reasoning-parser deepseek_r1)4 random seeds per benchmark,
temperature=1.0,top_p=1.0,max_output_tokens=65536.Repro: AIME24 (math_with_judge)
Requires AIME24 validation data pre-downloaded to the cluster (see data download step below).
Repro: LCB (code_gen)
Requires LiveCodeBench v5 validation data pre-downloaded to the cluster (see data download step below).
Data Download
Both datasets must be downloaded to the Gym data directories on the cluster before running rollouts. This uses
ng_download_dataset_from_hfinside the nemo-rl container:Summary by CodeRabbit
New Features
Bug Fixes
Chores
Refactor
Tests